// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <inttypes.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <threads.h>

#include <zircon/assert.h>
#include <zircon/process.h>
#include <zircon/status.h>
#include <zircon/syscalls.h>
#include <zircon/syscalls/exception.h>
#include <zircon/syscalls/port.h>
#include <zircon/threads.h>

#include <fbl/unique_ptr.h>
#include <fbl/vector.h>
#include <inspector/inspector.h>
#include <pretty/hexdump.h>
#include <task-utils/get.h>

static int verbosity_level = 0;

void print_error(const char* fmt, ...) {
  va_list args;
  va_start(args, fmt);
  fprintf(stderr, "ERROR: ");
  vfprintf(stderr, fmt, args);
  fprintf(stderr, "\n");
  va_end(args);
}

void print_zx_error(zx_status_t status, const char* fmt, ...) {
  va_list args;
  va_start(args, fmt);
  fprintf(stderr, "ERROR: ");
  vfprintf(stderr, fmt, args);
  fprintf(stderr, ": %d(%s)", status, zx_status_get_string(status));
  fprintf(stderr, "\n");
  va_end(args);
}

// While this should never fail given a valid handle,
// returns ZX_KOID_INVALID on failure.
zx_koid_t get_koid(zx_handle_t handle) {
  zx_info_handle_basic_t info;
  if (zx_object_get_info(handle, ZX_INFO_HANDLE_BASIC, &info, sizeof(info), NULL, NULL) < 0) {
    // This shouldn't ever happen, so don't just ignore it.
    print_error("Eh? ZX_INFO_HANDLE_BASIC failed");
    return ZX_KOID_INVALID;
  }
  return info.koid;
}

// How much memory to dump, in bytes.
// Space for this is allocated on the stack, so this can't be too large.
constexpr size_t kMemoryDumpSize = 256;

void dump_memory(zx_handle_t proc, uintptr_t start, size_t len) {
  // Make sure we're not allocating an excessive amount of stack.
  ZX_DEBUG_ASSERT(len <= kMemoryDumpSize);

  uint8_t buf[len];
  auto res = zx_process_read_memory(proc, start, buf, len, &len);
  if (res < 0) {
    printf("failed reading %p memory; error : %d\n", (void*)start, res);
  } else if (len != 0) {
    hexdump_ex(buf, len, start);
  }
}

void dump_thread(zx_handle_t process, inspector_dsoinfo_t* dso_list, uint64_t tid,
                 zx_handle_t thread) {
  zx_thread_state_general_regs_t regs;
  zx_vaddr_t pc = 0, sp = 0, fp = 0;

  if (inspector_read_general_regs(thread, &regs) != ZX_OK) {
    // Error message has already been printed.
    return;
  }

#if defined(__x86_64__)
  pc = regs.rip;
  sp = regs.rsp;
  fp = regs.rbp;
#elif defined(__aarch64__)
  pc = regs.pc;
  sp = regs.sp;
  fp = regs.r[29];
#else
  // It's unlikely we'll get here as trying to read the regs will likely
  // fail, but we don't assume that.
  printf("unsupported architecture .. coming soon.\n");
  return;
#endif

  char thread_name[ZX_MAX_NAME_LEN];
  auto status = zx_object_get_property(thread, ZX_PROP_NAME, thread_name, sizeof(thread_name));
  if (status < 0) {
    strlcpy(thread_name, "unknown", sizeof(thread_name));
  }

  printf("<== Thread %s[%" PRIu64 "] ==>\n", thread_name, tid);

  inspector_print_general_regs(stdout, &regs, nullptr);

  printf("bottom of user stack:\n");
  dump_memory(process, sp, kMemoryDumpSize);

  inspector_print_backtrace_markup(stdout, process, thread, dso_list, pc, sp, fp, true);

  if (verbosity_level >= 1)
    printf("Done handling thread %" PRIu64 ".%" PRIu64 ".\n", get_koid(process), get_koid(thread));
}

void dump_all_threads(uint64_t pid, zx_handle_t process) {
  // First get the thread count so that we can allocate an appropriately
  // sized buffer. This is racy but it's the nature of the beast.
  size_t num_threads;
  zx_status_t status =
      zx_object_get_info(process, ZX_INFO_PROCESS_THREADS, nullptr, 0, nullptr, &num_threads);
  if (status != ZX_OK) {
    print_zx_error(status, "failed to get process thread info (#threads)");
    exit(1);
  }

  auto threads = fbl::unique_ptr<zx_koid_t[]>(new zx_koid_t[num_threads]);
  size_t records_read;
  status = zx_object_get_info(process, ZX_INFO_PROCESS_THREADS, threads.get(),
                              num_threads * sizeof(threads[0]), &records_read, nullptr);
  if (status != ZX_OK) {
    print_zx_error(status, "failed to get process thread info");
    exit(1);
  }
  ZX_DEBUG_ASSERT(records_read == num_threads);

  printf("%zu thread(s)\n", num_threads);

  inspector_dsoinfo_t* dso_list = inspector_dso_fetch_list(process);
  inspector_print_markup_context(stdout, process);

  // TODO(dje): Move inspector's DebugInfoCache here, so that we can use it
  // across all threads.

  for (size_t i = 0; i < num_threads; ++i) {
    zx_koid_t tid = threads[i];
    zx_handle_t thread;
    // TODO(dje): There is value in specifying exactly the rights we need,
    // but an explicit list this early has a higher risk of bitrot.
    status = zx_object_get_child(process, tid, ZX_RIGHT_SAME_RIGHTS, &thread);
    if (status < 0) {
      printf("WARNING: failed to get a handle to [%" PRIu64 ".%" PRIu64 "] : error %d\n", pid, tid,
             status);
      continue;
    }

    zx_handle_t suspend_token = ZX_HANDLE_INVALID;
    status = zx_task_suspend_token(thread, &suspend_token);
    if (status != ZX_OK) {
      print_zx_error(status, "unable to suspend thread, skipping");
      zx_handle_close(thread);
      continue;
    }

    zx_signals_t observed = 0u;
    // Try to be robust and don't wait forever. The timeout is a little
    // high as we want to work well in really loaded systems.
    auto deadline = zx_deadline_after(ZX_SEC(5));
    // Currently, asking to wait for suspended means only waiting for the
    // thread to suspend. If the thread terminates instead this will wait
    // forever (or until the timeout). Thus we need to explicitly wait for
    // ZX_THREAD_TERMINATED too.
    zx_signals_t signals = ZX_THREAD_SUSPENDED | ZX_THREAD_TERMINATED;
    status = zx_object_wait_one(thread, signals, deadline, &observed);
    if (status == ZX_OK) {
      if (observed & ZX_THREAD_TERMINATED) {
        printf("Unable to print backtrace of thread %" PRIu64 ".%" PRIu64 ": terminated\n", pid,
               tid);
      } else {
        dump_thread(process, dso_list, tid, thread);
      }
    } else {
      print_zx_error(status,
                     "failure waiting for thread %" PRIu64 ".%" PRIu64 " to suspend, skipping", pid,
                     tid);
    }

    zx_handle_close(suspend_token);
    zx_handle_close(thread);
  }

  inspector_dso_free_list(dso_list);
}

void usage(FILE* f) {
  fprintf(f, "Usage: threads [options] pid\n");
  fprintf(f, "Options:\n");
  fprintf(f, "  -v[n] = set verbosity level to N\n");
}

int main(int argc, char** argv) {
  zx_status_t status;
  zx_koid_t pid = ZX_KOID_INVALID;

  int i;
  for (i = 1; i < argc; ++i) {
    const char* arg = argv[i];
    if (arg[0] == '-') {
      if (strcmp(arg, "-h") == 0 || strcmp(arg, "--help") == 0) {
        usage(stdout);
        return 0;
      } else if (strncmp(arg, "-v", 2) == 0) {
        if (arg[2] != '\0') {
          verbosity_level = atoi(arg + 2);
        } else {
          verbosity_level = 1;
        }
      } else {
        usage(stderr);
        return 1;
      }
    } else {
      break;
    }
  }
  if (i == argc || i + 1 != argc) {
    usage(stderr);
    return 1;
  }
  char* endptr;
  const char* pidstr = argv[i];
  pid = strtoull(pidstr, &endptr, 0);
  if (!(pidstr[0] != '\0' && *endptr == '\0')) {
    fprintf(stderr, "ERROR: invalid pid: %s", pidstr);
    return 1;
  }

  inspector_set_verbosity(verbosity_level);

  zx_handle_t thread_self = thrd_get_zx_handle(thrd_current());
  if (thread_self == ZX_HANDLE_INVALID) {
    print_error("unable to get thread self");
    return 1;
  }

  zx_handle_t process;
  zx_obj_type_t type;
  status = get_task_by_koid(pid, &type, &process);
  if (status < 0) {
    print_zx_error(status, "unable to get a handle to %" PRIu64, pid);
    return 1;
  }

  if (type != ZX_OBJ_TYPE_PROCESS) {
    print_error("PID %" PRIu64 " is not a process. Threads can only be dumped from processes", pid);
    return 1;
  }

  char process_name[ZX_MAX_NAME_LEN];
  status = zx_object_get_property(process, ZX_PROP_NAME, process_name, sizeof(process_name));
  if (status < 0) {
    strlcpy(process_name, "unknown", sizeof(process_name));
  }

  printf("Backtrace of threads of process %" PRIu64 ": %s\n", pid, process_name);

  dump_all_threads(pid, process);
  zx_handle_close(process);

  return 0;
}
